import java.util.Collection; 
  
  
/** 
 * CS 1332 Fall 2013 
 * AVL Tree 
 *  
 * In this class, you will program an AVL Tree (Adelson Veskii-Landis Tree). 
 * This is like a better version of a binary search tree in that it  
 * tries to fill out every level of the tree as much as possible. It 
 * accomplishes this by keeping track of each node's height and balance 
 * factor. As you recurse back up from operations that modify the tree 
 * (like add or remove), you will update the height and balance factor 
 * of the current node, and perform a rotation on the current node if  
 * necessary. Keeping this in mind, let's get started! 
 *  
 * **************************NOTE************************************* 
 * please please please  treat null as positive infinity!!!!!!!! 
 * PLEASE TREAT NULL AS POSITIVE INFINITY!!!! 
 * *************************NOTE************************************** 
 *  
 * I STRONLY RECOMMEND THAT YOU IMPLEMENT THIS DATA STRUCTURE RECURSIVELY! 
 *  
 * Please make any new internal classes, instance data, and methods private!! 
 *  
 * DO NOT CHANGE ANY OF THE PUBLIC METHOD HEADERS 
 */
public class AVL<T extends Comparable<T>> { 
  
    private AVLNode<T> root; 
    private int size; 
  
    /** 
     * I promise you, this is just like the add() method you coded 
     * in the BST part of the homework! You will start off at the 
     * root and find the proper place to add the data. As you  
     * recurse back up the tree, you will have to update the 
     * heights and balance factors of each node that you visited 
     * while reaching the proper place to add your data. Immediately 
     * before you return out of each recursive step, you should update 
     * the height and balance factor of the current node and then 
     * call rotate on the current node. You will then return the node 
     * that comes from the rotate(). This way, the re-balanced subtrees 
     * will properly be added back to the whole tree. Also, don't forget 
     * to update the size of the tree as a whole. 
     *  
     * PLEASE TREAT NULL AS POSITIVE INFINITY!!!! 
     *  
     * @param data The data do be added to the tree. 
     */
    public void add(T data) { 
        if (root==null) 
            root = new AVLNode<T>(data); 
        else if (data==null || root.getData()!=null && data.compareTo(root.getData())>0) 
            add(root.getRight(),data,root,true); 
        else add(root.getLeft(),data,root,false); 
        size++; 
          
        updateHeightAndBF(root);                   // important 
        rotate(root,null,true);} 
  
    private void add(AVLNode<T> node, T data, AVLNode<T> prev, boolean right) { 
  
        if (node==null) { 
            node = new AVLNode<T>(data); 
            if (right)  
                prev.setRight(node); 
            else
                prev.setLeft(node); } 
        else if (data==null || node.getData()!=null && data.compareTo(node.getData())>0) 
            add(node.getRight(),data,node,true);             
        else add(node.getLeft(),data,node,false);  
  
        updateHeightAndBF(node); 
        rotate(node,prev,right); 
    } 
  
  
  
    /** 
     * This is a pretty simple method. All you need to do is to get 
     * every element in the collection that is passed in into the tree. 
     *  
     * Try to think about how you can combine a for-each loop and your 
     * add method to accomplish this. 
     *  
     * @param c A collection of elements to be added to the tree. 
     */
    public void addAll(Collection<? extends T> c) { 
        for (T item: c) 
            add(item); } 
  
    /** 
     * All right, now for the remove method. Just like in the vanilla BST, you 
     * will have to traverse to find the data the user is trying to remove.  
     *  
     * You will have three cases: 
     *  
     * 1. Node to remove has zero children. 
     * 2. Node to remove has one child. 
     * 3. Node to remove has two children. 
     *  
     * For the first case, you simply return null up the tree. For the second case, 
     * you return the non-null child up the tree.  
     *  
     * Just as in add, you'll have to updateHeightAndBF() as well as rotate() just before 
     * you return out of each recursive step. 
     *  
     * FOR THE THIRD CASE USE THE PREDECESSOR OR YOU WILL LOSE POINTS 
     *  
     * @param data The data to search in the tree. 
     * @return The data that was removed from the tree. 
     */
    public T remove(T data) { 
        if (!contains(data)) 
            return null; 
  
        boolean removed = false; 
        if (root==null) return null; 
        if (data.equals(root.getData())) { 
            removed = true; 
            if (root.getLeft()==null && root.getRight()==null)  // no children 
                clear(); 
            else if (root.getLeft()!=null && root.getRight()!=null) 
                root.setData(remove(findMax(root.getLeft()).getData())); 
            else if (root.getLeft()!=null && root.getRight()==null) 
                root = root.getLeft(); 
            else 
                root = root.getRight(); } 
        else if (data==null || root.getData()!=null && data.compareTo(root.getData())>0) 
            removed =removeChild(root.getRight(), data, root, true); 
        else removed =removeChild(root.getLeft(), data, root, false); 
          
        if (removed) { 
            updateHeightAndBF(root); 
            rotate(root,null,true);                   // important 
            size--; 
            return data; } 
        else return null; 
    } 
  
  
    private boolean removeChild(AVLNode<T> node, T data, AVLNode<T> prev, boolean right) { 
        if (node==null) return false; 
        if (data.equals(node.getData())) 
            if (node.getLeft()==null && node.getRight()==null)  // no children 
                if (right) prev.setRight(null); 
                else prev.setLeft(null); 
            else if (node.getLeft()!=null && node.getRight()!=null) 
                node.setData(remove(findMax(node.getLeft()).getData())); 
            else if (node.getLeft()!=null && node.getRight()==null) 
                if (right) 
                    prev.setRight(node.getLeft()); 
                else
                    prev.setLeft(node.getLeft()); 
            else 
                if (right) 
                    prev.setRight(node.getRight()); 
                else
                    prev.setLeft(node.getRight()); 
        else if (data==null || node.getData()!=null && data.compareTo(node.getData())>0) 
            removeChild(node.getRight(), data, node, true); 
        else removeChild(node.getLeft(), data, node, false); 
  
        updateHeightAndBF(node); 
        rotate(node,prev,right); 
  
        return true; 
    } 
  
    private AVLNode<T> findMax(AVLNode<T> node) { 
        if (node.getRight()==null) 
            return node; 
        else return findMax(node.getRight()); 
    } 
  
    /** 
     * This method should be pretty simple, all you have to do is recurse 
     * to the left or to the right and see if the tree contains the data. 
     *  
     * @param data The data to search for in the tree. 
     * @return The boolean flag that indicates if the data was found in the tree or not. 
     */
    public boolean contains(T data) {        
        return contains(root,data); } 
  
    private boolean contains(AVLNode<T> node, T data) { 
        if (node==null) 
            return false;    
        else if (node.getData()==null && data==null) 
            return true; 
        else if (data!=null && data.equals(node.getData())) 
            return true; 
        else if (data==null || node.getData()!=null && data.compareTo(node.getData())>0) 
            return contains(node.getRight(),data);           
        else 
            return contains(node.getLeft(),data);  
    } 
  
  
    /** 
     * Again, simply recurse through the tree and find the data that is passed in. 
     *  
     * @param data The data to fetch from the tree. 
     * @return The data that the user wants from the tree. Return null if not found. 
     */
    public T get(T data) { 
        if (contains(data)) return data; 
        else return null; 
    } 
  
  
    /** 
     * Test to see if the tree is empty. 
     *  
     * @return A boolean flag that is true if the tree is empty. 
     */
    public boolean isEmpty() { 
        return root==null; } 
  
    /** 
     * Return the number of data in the tree. 
     *  
     * @return The number of data in the tree. 
     */
    public int size() { 
        return size; } 
  
    /** 
     * Reset the tree to its original state. Get rid of every element in the tree. 
     */
    public void clear() { 
        size = 0; 
        root = null; } 
  
    // The below methods are all private, so we will not be directly grading them, 
    // however we strongly recommend you not change them, and make use of them. 
  
  
    /** 
     * Use this method to update the height and balance factor for a node. 
     *  
     * @param node The node whose height and balance factor need to be updated. 
     */
    private void updateHeightAndBF(AVLNode<T> node) { 
            if (node!=null) { 
                if (node.getLeft()==null && node.getRight()==null) {  // no children 
                    node.setHeight(1);                   // important 
                    node.setBF(0); } 
                else if (node.getLeft()!=null && node.getRight()!=null) { 
                    node.setHeight(Math.max(node.getLeft().getHeight(), node.getRight().getHeight())+1); 
                    node.setBF(node.getLeft().getHeight() - node.getRight().getHeight()); } 
                else if (node.getLeft()!=null && node.getRight()==null) { 
                    node.setHeight(node.getLeft().getHeight()+1); 
                    node.setBF(node.getLeft().getHeight()); } 
                else { 
                    node.setHeight(node.getRight().getHeight()+1); 
                    node.setBF(-node.getRight().getHeight()); } 
            } 
        } 
  
        /** 
         * In this method, you will check the balance factor of the node that is passed in and 
         * decide whether or not to perform a rotation. If you need to perform a rotation, simply 
         * call the rotation and return the new root of the balanced subtree. If there is no need 
         * for a rotation, simply return the node that was passed in. 
         *  
         * @param node - a potentially unbalanced node 
         * @return The new root of the balanced subtree. 
         */
        private AVLNode<T> rotate(AVLNode<T> node, AVLNode<T> prev, boolean right) { 
            if (node.getBf()>1) { 
                if (node.getLeft().getBf()==1) // right 
                    return rightRotate(node,prev,right); 
                else if (node.getLeft().getBf()==-1) // left-right 
                    return leftRightRotate(node,prev,right);} 
            else if (node.getBf()<-1) { 
                if (node.getRight().getBf()==1) // right-left 
                    return rightLeftRotate(node,prev,right); 
                else if (node.getRight().getBf()==-1) // left 
                    return leftRotate(node,prev,right); } 
            return node; } 
      
      
  
        /** 
         * In this method, you will perform a left rotation. Remember, you perform a  
         * LEFT rotation when the sub-tree is RIGHT heavy. This moves more nodes over to 
         * the LEFT side of the node that is passed in so that the height differences 
         * between the LEFT and RIGHT subtrees differ by at most one. 
         *  
         * HINT: DO NOT FORGET TO RE-CALCULATE THE HEIGHT OF THE NODES 
         * WHOSE CHILDREN HAVE CHANGED! YES, THIS DOES MAKE A DIFFERENCE! 
         *  
         * @param node - the current root of the subtree to rotate. 
         * @return The new root of the subtree 
         */
        private AVLNode<T> leftRotate(AVLNode<T> node, AVLNode<T> prev, boolean right) { 
            AVLNode<T> newRoot = node.getRight(); 
            if (prev==null) {                   // important (dealing with root) 
                root = newRoot; 
                node.setRight(newRoot.getLeft()); 
                newRoot.setLeft(node);  
                return newRoot; } 
            if (right) { 
                prev.setRight(newRoot); 
                node.setRight(newRoot.getLeft()); 
                newRoot.setLeft(node); } 
            else { 
                prev.setLeft(newRoot); 
                node.setRight(newRoot.getLeft()); 
                newRoot.setLeft(node); } 
            updateHeightAndBF(node);                       // important 
            updateHeightAndBF(newRoot); 
            updateHeightAndBF(prev); 
            return newRoot; 
        } 
  
        /** 
         * In this method, you will perform a right rotation. Remember, you perform a 
         * RIGHT rotation when the sub-tree is LEFT heavy. THis moves more nodes over to 
         * the RIGHT side of the node that is passed in so that the height differences 
         * between the LEFT and RIGHT subtrees differ by at most one. 
         *  
         * HINT: DO NOT FORGET TO RE-CALCULATE THE HEIGHT OF THE NODES 
         * WHOSE CHILDREN HAVE CHANGED! YES, THIS DOES MAKE A DIFFERENCE! 
         *  
         * @param node - The current root of the subtree to rotate. 
         * @return The new root of the rotated subtree. 
         */
        private AVLNode<T> rightRotate(AVLNode<T> node, AVLNode<T> prev, boolean right) { 
            AVLNode<T> newRoot = node.getLeft(); 
            if (prev==null) { 
                root = newRoot; 
                node.setLeft(newRoot.getRight()); 
                newRoot.setRight(node);  
                return newRoot; } 
            if (right) { 
                prev.setRight(newRoot); 
                node.setLeft(newRoot.getRight()); 
                newRoot.setRight(node); } 
            else { 
                prev.setLeft(newRoot); 
                node.setLeft(newRoot.getRight()); 
                newRoot.setRight(node); } 
            updateHeightAndBF(node); 
            updateHeightAndBF(newRoot); 
            updateHeightAndBF(prev); 
            return newRoot; 
        } 
  
        /** 
         * In this method, you will perform a left-right rotation. You can simply use 
         * the left and right rotation methods on the node and the node's child. Remember 
         * that you must perform the rotation on the node's child first, otherwise you will 
         * end up with a mangled tree (sad face). After rotating the child, remember to link up 
         * the new root of the that first rotation with the node that was passed in. 
         *  
         * The whole point of heterogeneous rotations is to transform the node's  
         * subtree into one of the cases handled by the left and right rotations. 
         *  
         * @param node 
         * @return The new root of the subtree. 
         */
        private AVLNode<T> leftRightRotate(AVLNode<T> node, AVLNode<T> prev, boolean right) { 
            if (prev==null) { 
                leftRotate(node.getLeft(),node,false); 
                AVLNode<T> newRoot = rightRotate(node,prev,right); 
                updateHeightAndBF(node); 
                updateHeightAndBF(newRoot); 
                return newRoot; } 
            leftRotate(node.getLeft(),node,false); 
            AVLNode<T> newRoot = rightRotate(node,prev,right); 
            updateHeightAndBF(node); 
            updateHeightAndBF(newRoot); 
            updateHeightAndBF(prev); 
            return newRoot; 
        } 
  
        /** 
         * In this method, you will perform a right-left rotation. You can simply use your 
         * right and left rotation methods on the node and the node's child. Remember 
         * that you must perform the rotation on the node's child first, otherwise 
         * you will end up with a mangled tree (super sad face). After rotating the node's child, 
         * remember to link up the new root of that first rotation with the node that was 
         * passed in. 
         *  
         * Again, the whole point of the heterogeneous rotations is to first transform the 
         * node's subtree into one of the cases handled by the left and right rotations. 
         *  
         * @param node 
         * @return The new root of the subtree. 
         */
        private AVLNode<T> rightLeftRotate(AVLNode<T> node, AVLNode<T> prev, boolean right) { 
            if (prev==null) { 
                rightRotate(node.getRight(),node,true); 
                AVLNode<T> newRoot = leftRotate(node,prev,right); 
                updateHeightAndBF(node); 
                updateHeightAndBF(newRoot); 
                return newRoot; } 
            rightRotate(node.getRight(),node,true); 
            AVLNode<T> newRoot = leftRotate(node,prev,right); 
            updateHeightAndBF(node); 
            updateHeightAndBF(newRoot); 
            updateHeightAndBF(prev); 
            return newRoot; 
        } 
  
        public AVLNode<T> getRoot(){ //added by author 
            return root;} 
  
        public void setRoot(AVLNode<T> n1) { //added by author 
            root = n1;} 
  
        public void setSize(int i) { //added by author 
            size =  i;} 
  
    } 